// SPDX-License-Identifier: BSD-2-Clause
/*
 * Copyright (c) 2018, Linaro Limited
 */
#include <pthread.h>
#include "socket_zmq.h"
#include "selfqueue.h"
#include <stdlib.h>
#include "qrng.h"
#include <dlfcn.h>
pthread_cond_t cond;
pthread_mutex_t mutex;

#define RANDOMBYTES 1024
PSKF_FUNCLIST FunctionList;
void qrng_test(){
	int ret = 0;
	char devices[128] = {0};
	u32 devices_size = sizeof(devices);
	HANDLE hdev;

	ret = load_library();
	if (ret) {
		printf("load_library() failed: %#x\n", ret);
		return ret;
	}

	ret = FunctionList->SKF_EnumDev(1, devices, &devices_size);
	if (ret) {
		printf("SKF_EnumDev() failed: %#x\n", ret);
		goto end;
	}

	ret = FunctionList->SKF_ConnectDev(devices, &hdev);
	if (ret) {
		printf("SKF_ConnectDev() failed: %#x\n", ret);
		goto end;
	}

    ret = get_device_info(hdev);
    if (ret) {
		printf("get device info failed: %#x\n", ret);
        goto end;
    }

	unsigned char random_data[RANDOMBYTES] = {0};

	ret = FunctionList->SKF_GenRandom(hdev, random_data, RANDOMBYTES);
	if (ret) {
		printf("sym_test() failed: %#x\n", ret);
	}
	
	print_data("Random Data:", random_data, 256);
	end:
	if (hdev) ret = FunctionList->SKF_DisConnectDev(hdev);
	// return 0;
}

typedef struct identify_params
{
	ca_oppo oparams;
	ca_self sparams;
	TEEC_Context ctx_1;
	TEEC_Session sess_1;
} iden_params;

typedef struct socket_params
{
	ca_oppo *oparams;
	ca_self *sparams;
	TEEC_Context ctx_1;
	TEEC_Session sess_1;
} socket_params;

// node_data:
char local_ipset[] = "tcp://*:5566";
char local_oppo_ipset[] = "tcp://localhost:5577";

char *nodeID;
char *nodeIP;
// 1.身份认证 ; 2.内容认证
TEEC_UUID uuid1 = TA_LIBS_EXAMPLE_UUID;
TEEC_UUID uuid2 = TA2_LIBS_EXAMPLE_UUID;
TEEC_Context ctx1, ctx2;
TEEC_Session sess1, sess2;
TEEC_Result res;
uint32_t err_origin;

//开启TA1 TA2
void start_ta1_ta2()
{
	res = TEEC_InitializeContext(NULL, &ctx1);
	if (res != TEEC_SUCCESS)
		errx(1, "TEEC_InitializeContext failed with code 0x%x", res);
	res = TEEC_OpenSession(&ctx1, &sess1, &uuid1,
						   TEEC_LOGIN_PUBLIC, NULL, NULL, &err_origin);
	if (res != TEEC_SUCCESS)
		errx(1, "TEEC_Opensession unexpectedly returned with code 0x%x"
				" origin 0x%x",
			 res, err_origin);
	res = TEEC_InitializeContext(NULL, &ctx2);
	if (res != TEEC_SUCCESS)
		errx(1, "TEEC_InitializeContext failed with code 0x%x", res);
	res = TEEC_OpenSession(&ctx2, &sess2, &uuid2,
						   TEEC_LOGIN_PUBLIC, NULL, NULL, &err_origin);
	if (res != TEEC_SUCCESS)
		errx(1, "TEEC_Opensession unexpectedly returned with code 0x%x"
				" origin 0x%x",
			 res, err_origin);
}
//当前节点数据初始化
ca_self self1;
ca_oppo oppo1;

// node端服务器
void *node_server(socket_params *sock_par)
{
	zmq_Server(local_ipset, ctx1, sess1, ctx2, sess2, sock_par->oparams, sock_par->sparams);
}

void *pqc_kem_func(char *opponode_ipset)
{
	printf("\n**************************************************************\n");
	printf("                  Identity Authentication\n");
	ca_self self;
	ca_oppo oppo;
	int pwd = 1;
	while(pwd < 3)
	{
		// printf("pwd=%d\n", pwd);
		char send_buffer[3296] = {0};
		int send_buffer_len = 3295;
		char rec_buffer[3296] = {0};
		pqc_node_client_send_handler(pwd, send_buffer);
		// printf("ok\n");
		int rec_val = zmq_Client(opponode_ipset, send_buffer, send_buffer_len, rec_buffer);
		// printf("rec_val:%d\n",rec_val);
		if (rec_val != 0)
		{
			continue;
		}
		node_client_receive_handler(rec_buffer, ctx1, sess1, ctx2, sess2, &self, &oppo);
		pwd++;
	}
	if(pwd == 3){
		Sk_Secure_Read(ctx2, sess2, oppo.oppoid);
		// printf("Secure Read Finish\n");
	}

	// verify key
	char send_buffer[3296] = {0};
	int send_buffer_len = 3295;
	char rec_buffer[3296] = {0};
	pqc_node_client_send_handler(pwd, send_buffer);
	// printf("ok\n");
	int rec_val = zmq_Client(opponode_ipset, send_buffer, send_buffer_len, rec_buffer);
	// printf("rec_val:%d\n",rec_val);
	node_client_receive_handler(rec_buffer, ctx1, sess1, ctx2, sess2, &self, &oppo);

	//save opponode_information
	if (hasheadnode == 0)
	{
		node_data_init(&head_node, oppo.oppoid, opponode_ipset);
		hasheadnode = 1;
	}
	else
	{
		node_data_add(&head_node, oppo.oppoid, opponode_ipset);
	}
	apply_key_flag = 1;
	printf("\nIdentification Authentication is \033[1;32mAccomplished\033[0m\n");
	printf("**************************************************************\n");
}

// node与node交互 内容认证
void *node_context(char *opponode_ipset)
{
	/// MAIN FUNC
	printf("\nStart Context Authentication\n");
	ca_self self;
	ca_oppo oppo;
	
	//身份认证密钥数据复制到局部变量中
	params_init(&oppo, &self, nodeID, nodeIP);

	// IP+ID结构体构造
	uint8_t context_oppoid[8];
	uint8_t node_new_sk_flag = 0;
	if (!node_data_search(&head_node, opponode_ipset, context_oppoid, &node_new_sk_flag))
	{
		printf("\nNot found its ID!\n");
		return;
	}
	//memmove(oppo.oppoid, oppo1.oppoid, 8);
	int pwd = 6;

	//循环处理输入数据
	struct timeval start, end;
	char send_buffer[3296] = {0};
	int send_buffer_len = 3295;
	char rec_buffer[3296] = {0};
	for (int i = 1; i < 5; ++i)
	{
		//printf("\n%d, %d\n", sizeof(data), sizeof(data[0]));
		printf("No.%d Content Authenticaton\n", i);
		char data[32] = {0};
		time_t timestamp;
		time(&timestamp);
		// sprintf(data, "%d", timestamp);
		struct tm tm;
		tm = *localtime(&timestamp);
		strftime(data, sizeof(data), "%Y%m%d%H%M%S", &tm);

		memmove(data+strlen(data), ",", 1);
		generate_string(data+strlen(data), 8);
		rec_buffer[0] = '0';
		self.data = data;
		printf("\nOriginal data:");
		printf("\n\033[1;34m%s\033[0m\n", self.data);
		gettimeofday(&start, NULL);
		Context_Hmac(ctx2, sess2, context_oppoid, self.data, self.datapackage, &node_new_sk_flag);
		
		node_client_send_handler(pwd, send_buffer, &self, &node_new_sk_flag);
		int rec_val = zmq_Client(opponode_ipset, send_buffer, send_buffer_len, rec_buffer);
		if (rec_val != 0)
		{
			//i--;
			printf("\nCommunication_Wrong!\n");
			continue;
		}
		node_client_receive_handler(rec_buffer, ctx1, sess1, ctx2, sess2, &self, &oppo);

		gettimeofday(&end, NULL);
		if(node_new_sk_flag == 1 && self.con_flag[0] == 0x01){
			node_new_sk_flag = 0;
			node_data_update_skflag(&head_node, opponode_ipset);
		}
		long timeuse = 1000000 * (end.tv_sec - start.tv_sec) + end.tv_usec - start.tv_usec;
		printf("Content Authentication Time=%.3fms\n", timeuse / 1000.0);
		printf("\nNo.%d Content Authentication \033[1;32mfinished\033[0m\n", i);
		printf("-------------------------------------------------------\n");
		sleep(1);
	}
}

void *node_context_attack(char *opponode_ipset)
{
	/// MAIN FUNC
	printf("\nAttack test start:\n");
	ca_self self;
	ca_oppo oppo;
	
	//身份认证密钥数据复制到局部变量中
	params_init(&oppo, &self, nodeID, nodeIP);

	// IP+ID结构体构造
	uint8_t context_oppoid[8];
	uint8_t sk_update_flag = 0;
	if (!node_data_search(&head_node, opponode_ipset, context_oppoid, &sk_update_flag))
	{
		printf("\nNot found its ID!\n");
		return;
	}
	// memmove(oppo.oppoid, oppo1.oppoid, 8);
	int pwd = 6;
	uint8_t lastpackage[96] = {0};
	uint8_t node_new_sk_flag = 1;
	//循环处理输入数据
	char *attack_test_data1 = {"20231106181715,test1"};
	char *attack_test_data2 = {"20231106181716,test2"};
	for (int i = 0; i < 2; ++i)
	{
		char send_buffer[3350] = {0};
		int send_buffer_len = 3349;
		char rec_buffer[3350] = {0};
		if (i == 0)
		{
			self.data = attack_test_data1;
		}	
		else
		{
			self.data = attack_test_data2;
		}
		Context_Hmac(ctx2, sess2, context_oppoid, self.data, self.datapackage, &node_new_sk_flag);
		printf("\nNo.%d content authentication\nOriginal Datapackage:\n", i + 1);
		format_print(self.datapackage, 96);
		//篡改攻击 + 重放攻击
		int attack_id = 0;
		if (i == 0)
		{
			printf("\nORDER:1.Tampering Attack 2.Replay Attack\nPlease Input Your Order:");
			scanf("%d", &attack_id);
			getchar();
			switch (attack_id)
			{
			case 1:
				printf("\nNo.%d Content Authentication\nOriginal Datapackage:\n", i + 1);
				format_print(self.datapackage, 96);
				for (int j = 0; j < 2; ++j)
				{
					self.datapackage[j] = 0x10;
				}	
				printf("\nTampering Attack Changed Datapackage:\n");
				format_print(self.datapackage, 96);
				i += 2;
				break;
			case 2:
				memmove(lastpackage, self.datapackage, 96);
				break;
			default:
				break;
			}
		}
		if (i < 2)
		{
			memmove(self.datapackage, lastpackage, 96);
			printf("\nReplay Attack Changed Datapackage:\n");
			format_print(self.datapackage, 96);
		}
		// printf("node_new_sk_flag:%d\n",node_new_sk_flag);
		// printf("sk_update_flag:%d\n",sk_update_flag);
		node_client_send_handler(pwd, send_buffer, &self, &node_new_sk_flag);
		int rec_val = zmq_Client(opponode_ipset, send_buffer, send_buffer_len, rec_buffer);
		if (rec_val != 0)
		{
			// i--;
			printf("\nCommunication_Wrong!\n");
			continue;
		}
		node_client_receive_handler(rec_buffer, ctx1, sess1, ctx2, sess2, &self, &oppo);
		node_new_sk_flag = 0;
		sleep(1);
	}
	printf("\nAttack test end!\n");
}

void end_ta1_ta2()
{
	TEEC_CloseSession(&sess1);
	TEEC_FinalizeContext(&ctx1);
	TEEC_CloseSession(&sess2);
	TEEC_FinalizeContext(&ctx2);
}
/// TIME MODULE OVER

int main(int argc, char *argv[])
{
	printf("**************************************************************\n");
	printf("\033[1;32m");
	printf("                       LIPQUAP                               \n");
	printf("LIghtweight Post-QUantum cryptography Authentication Protocol\n");
	printf("                  Authentication Node                        \n");
	printf("\033[0m");
	printf("**************************************************************\n");
	// printf("Version:23.11.6v1\n");
	printf("Start TA\n");
	start_ta1_ta2();
	printf("Start TA is \033[1;32mAccomplished\033[0m\n");
	int ret;
	int button = 0;
	hasheadnode = 0;
	char *oppo_ipset = NULL;
	char nodeIDD[20];
	// char tempRID[8];
	char tempNodeID[9] = {0};
	generate_string(tempNodeID, 8);
	nodeID = tempNodeID;
	// nodeRID = tempNodeID;
	printf("\nNodeID:%s\n", nodeID);
	// format_print(nodeID, 8);
	memmove(Ca_Self_ID, nodeID, 8);
	// printf("Ca_Self_ID:%s\n", Ca_Self_ID);

	char local_IP[27] = {0};
	get_local_IP(local_IP);
	// printf("%s\n", local_IP);
	nodeIP = local_IP;
	// printf("%d\n", strcmp(nodeIP, node1_ipset));
	// printf("%d, %d", strlen(nodeIP), strlen(node1_ipset));
	printf("Node IP:%s\n",nodeIP);

	char oppo_IP[27] = {0};
	char input_oppo_IP[27] = {0};
	while(!checkIP(input_oppo_IP))
	{
		printf("Please Input IP Address of the Content Authentication Node:");
		gets(input_oppo_IP);
	}
	// printf("%s\n%d\n", input_oppo_IP, strlen(input_oppo_IP));
	memmove(oppo_IP, "tcp://", 6);
	memmove(oppo_IP + 6, input_oppo_IP, strlen(input_oppo_IP));
	memmove(oppo_IP + 6 + strlen(input_oppo_IP), ":5566", 5);
	oppo_ipset = oppo_IP;
	printf("Content Authentication Node IP:%s\n",oppo_ipset);

	pthread_t pqc_kem;
	params_init(&oppo1, &self1, nodeID, nodeIP);
	pthread_t apply_pk, node_ser;
	pthread_t ex_sk1, ex_sk2;
	pthread_t node_con1, node_con2;
	pthread_t attack_test;
	socket_params sock_par;
	sock_par.oparams = &oppo1;
	sock_par.sparams = &self1;
	
	res = sig_keypair_invoke(ctx1, sess1, Ca_Self_Sig_Pk);
	//节点开server
	ret = pthread_create(&node_ser, NULL, node_server, &sock_par);
	if (ret != 0)
	{
		printf("Create pthread error!\n");
		return 1;
	}

	while (1)
	{
		printf("\n**************************************************************\n");
		printf("Function:\n");
		printf("\t1.KEM; \n\t2.Node_Context; \n\t3.Attack_Test;\n\t0.Exit\n");
		printf("**************************************************************\n");
		printf("Please Input Your Order:");
		char c[10] = {0};
		gets(c);
		switch (c[0])
		{
		case '1': //pqc authen kem
			ret = pthread_create(&pqc_kem, NULL, pqc_kem_func, oppo_ipset);
			if (ret != 0)
			{
				printf("Create pqc_kem_pthread error!\n");
				return 1;
			}
			pthread_join(pqc_kem, NULL);
			break;
		case '2': //内容认证
			if (!apply_key_flag)
			{
				printf("Please run KEM first\n");
				break;
			}
			ret = pthread_create(&node_con1, NULL, node_context, oppo_ipset);
			if (ret != 0)
			{
				printf("Create node_context_pthread error!\n");
				return 1;
			}
			pthread_join(node_con1, NULL);
			break;
		case '3': // attack_test
			ret = pthread_create(&attack_test, NULL, node_context_attack, oppo_ipset);
			if (ret != 0)
			{
				printf("Create node_context_pthread error!\n");
				return 1;
			}
			pthread_join(attack_test, NULL);
			break;
		case '0':
			exit(0);
			break;
		default:
			break;
		}
	}
	pthread_join(node_ser, NULL);
	end_ta1_ta2();
}
